import numpy as np
import torch
import torch.nn.functional as F

from attack.base_attack import BaseAttack

class Attack(BaseAttack):
    def __init__(self, model, config):
        super().__init__('R+FGSM', model, config, batch_size = 100, targeted = False, llc = False)

    def attack_batch(self, images, labels):
        images = images + (self.config['alpha'] * self.config['epsilon'] * np.sign(np.random.randn(*images.shape))) #生成随机矩阵(注意是ndarray形式)
        images = np.clip(images, 0.0, 1.0).astype(np.float32)

        var_images = torch.from_numpy(images).to(self.model.device)
        var_images.requires_grad = True
        var_labels = torch.from_numpy(labels).to(self.model.device)

        eps = (1.0 - self.config['alpha']) * self.config['epsilon']                       #参数设定

        self.model.eval()
        output = self.model(var_images)
        loss = F.cross_entropy(output, var_labels)
        loss.backward()
        grad_sign = var_images.grad.sign().cpu().numpy()

        adv_images = images + eps * grad_sign
        adv_images = np.clip(adv_images, 0.0, 1.0)

        return adv_images